#!/usr/bin/env python3
"""
完整的D90标签剔除流程
"""

import os
from pathlib import Path

def remove_d90_labels():
    """剔除所有标签文件中的D90标签（类别9）"""
    base_path = Path("/root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset")
    
    total_files = 0
    modified_files = 0
    removed_labels = 0
    
    print("开始剔除D90标签...")
    
    # 处理训练集和验证集
    for subset in ['train', 'val']:
        labels_dir = base_path / "labels" / subset
        
        if not labels_dir.exists():
            print(f"标签目录不存在: {labels_dir}")
            continue
            
        print(f"\n处理 {subset} 集...")
        
        # 遍历所有标签文件
        for label_file in labels_dir.glob("*.txt"):
            total_files += 1
            
            try:
                # 读取原始标签
                with open(label_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                
                # 过滤掉D90标签（类别9）
                filtered_lines = []
                file_modified = False
                
                for line in lines:
                    line = line.strip()
                    if line:
                        parts = line.split()
                        if parts and parts[0] == '9':  # D90标签
                            removed_labels += 1
                            file_modified = True
                        else:
                            filtered_lines.append(line)
                
                # 如果文件被修改，写回文件
                if file_modified:
                    modified_files += 1
                    with open(label_file, 'w', encoding='utf-8') as f:
                        for line in filtered_lines:
                            f.write(line + '\n')
                    
            except Exception as e:
                print(f"    错误处理文件 {label_file}: {e}")
    
    print(f"\n标签剔除完成!")
    print(f"总共检查了 {total_files} 个标签文件")
    print(f"修改了 {modified_files} 个文件")
    print(f"剔除了 {removed_labels} 个D90标签")
    
    return total_files, modified_files, removed_labels

def update_data_configs():
    """更新数据配置文件，移除D90类别"""
    print("\n开始更新配置文件...")
    
    # 更新配置文件
    config_file = Path("/root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset/data.yaml")
    new_content = """train: /root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset/images/train
val: /root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset/images/val

nc: 9
names: ['D00', 'D10', 'D20', 'D30', 'D40', 'D50', 'D60', 'D70', 'D80']
"""
    with open(config_file, 'w', encoding='utf-8') as f:
        f.write(new_content)
    print(f"已更新: {config_file}")
    
    print("配置文件更新完成!")

def main():
    """主函数"""
    print("=" * 50)
    print("N-RDD2024数据集D90标签剔除工具")
    print("=" * 50)
    
    # 检查数据集目录是否存在
    if not Path("/root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset").exists():
        print("错误: 未找到merged_dataset目录")
        return
    
    # 剔除D90标签
    total_files, modified_files, removed_labels = remove_d90_labels()
    
    # 更新配置文件
    update_data_configs()
    
    print("\n" + "=" * 50)
    print("D90标签剔除完成!")
    print(f"总结:")
    print(f"- 检查了 {total_files} 个标签文件")
    print(f"- 修改了 {modified_files} 个文件")
    print(f"- 剔除了 {removed_labels} 个D90标签")
    print(f"- 类别数从10减少到9")
    print("=" * 50)

if __name__ == "__main__":
    main()